/* Copyright (c) 2023 OceanBase and/or its affiliates. All rights reserved.
miniob is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
         http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details. */

//
// Created by Wangyunlai on 2023/08/16.
//

#include "sql/optimizer/logical_plan_generator.h"

#include <algorithm>
#include <common/log/log.h>
#include <memory>
#include <utility>

#include "sql/expr/expression.h"
#include "sql/operator/calc_logical_operator.h"
#include "sql/operator/delete_logical_operator.h"
#include "sql/operator/explain_logical_operator.h"
#include "sql/operator/insert_logical_operator.h"
#include "sql/operator/update_logical_operator.h"
#include "sql/operator/join_logical_operator.h"
#include "sql/operator/logical_operator.h"
#include "sql/operator/predicate_logical_operator.h"
#include "sql/operator/project_logical_operator.h"
#include "sql/operator/table_get_logical_operator.h"
#include "sql/operator/group_by_logical_operator.h"

#include "sql/parser/parse_defs.h"
#include "sql/stmt/calc_stmt.h"
#include "sql/stmt/delete_stmt.h"
#include "sql/stmt/explain_stmt.h"
#include "sql/stmt/filter_stmt.h"
#include "sql/stmt/insert_stmt.h"
#include "sql/stmt/update_stmt.h"
#include "sql/stmt/select_stmt.h"
#include "sql/stmt/stmt.h"

#include "sql/expr/expression_iterator.h"
#include "storage/field/field_meta.h"

using namespace std;
using namespace common;

RC LogicalPlanGenerator::create(Stmt *stmt, unique_ptr<LogicalOperator> &logical_operator)
{
  RC rc = RC::SUCCESS;
  switch (stmt->type()) {
    case StmtType::CALC: {
      CalcStmt *calc_stmt = static_cast<CalcStmt *>(stmt);

      rc = create_plan(calc_stmt, logical_operator);
    } break;

    case StmtType::SELECT: {
      SelectStmt *select_stmt = static_cast<SelectStmt *>(stmt);

      rc = create_plan(select_stmt, logical_operator);
    } break;

    case StmtType::INSERT: {
      InsertStmt *insert_stmt = static_cast<InsertStmt *>(stmt);

      rc = create_plan(insert_stmt, logical_operator);
    } break;

    case StmtType::UPDATE: {
      UpdateStmt *update_stmt = static_cast<UpdateStmt *>(stmt);

      rc = create_plan(update_stmt, logical_operator);
    } break;

    case StmtType::DELETE: {
      DeleteStmt *delete_stmt = static_cast<DeleteStmt *>(stmt);

      rc = create_plan(delete_stmt, logical_operator);
    } break;

    case StmtType::EXPLAIN: {
      ExplainStmt *explain_stmt = static_cast<ExplainStmt *>(stmt);

      rc = create_plan(explain_stmt, logical_operator);
    } break;
    default: {
      rc = RC::UNIMPLEMENTED;
    }
  }
  return rc;
}

RC LogicalPlanGenerator::create_plan(CalcStmt *calc_stmt, std::unique_ptr<LogicalOperator> &logical_operator)
{
  logical_operator.reset(new CalcLogicalOperator(std::move(calc_stmt->expressions())));
  return RC::SUCCESS;
}

RC LogicalPlanGenerator::create_plan(SelectStmt *select_stmt, unique_ptr<LogicalOperator> &logical_operator)
{
  unique_ptr<LogicalOperator> *last_oper = nullptr;

  unique_ptr<LogicalOperator> table_oper(nullptr);
  last_oper = &table_oper;

  const std::vector<Table *> &tables = select_stmt->tables();

  // 只有一个表，不需要join
  if (tables.size() == 1) {
    unique_ptr<LogicalOperator> table_get_oper(new TableGetLogicalOperator(tables[0], ReadWriteMode::READ_ONLY));
    table_oper = std::move(table_get_oper);
  }
  // 两个以上的表，需要join
  else if (tables.size() > 1) {
    JoinLogicalOperator *join_oper = new JoinLogicalOperator;
    for (Table *table : tables) {
      unique_ptr<LogicalOperator> table_get_oper(new TableGetLogicalOperator(table, ReadWriteMode::READ_ONLY));
      join_oper->add_child(std::move(table_get_oper));
    }
    table_oper = unique_ptr<LogicalOperator>(join_oper);
  }
  // 0个表，输入非法
  else {
    return RC::INVALID_ARGUMENT;
  }

  // // 对于tables中的每一个表，遍历一遍
  // for (Table *table : tables) {

  //   unique_ptr<LogicalOperator> table_get_oper(new TableGetLogicalOperator(table, ReadWriteMode::READ_ONLY));
  //   if (table_oper == nullptr) {
  //     table_oper = std::move(table_get_oper);
  //     // 如果表多于一个，后边的表就需要join，这时候table_oper是第一个表,table_get_oper是第二个表
  //   } else {
  //     JoinLogicalOperator *join_oper = new JoinLogicalOperator;
  //     join_oper->add_child(std::move(table_oper));
  //     join_oper->add_child(std::move(table_get_oper));
  //     table_oper = unique_ptr<LogicalOperator>(join_oper);
  //   }
  // }

  unique_ptr<LogicalOperator> predicate_oper;

  RC rc = create_plan(select_stmt->filter_stmt(), predicate_oper);
  if (OB_FAIL(rc)) {
    LOG_WARN("failed to create predicate logical plan. rc=%s", strrc(rc));
    return rc;
  }

  if (predicate_oper) {
    if (*last_oper) {
      predicate_oper->add_child(std::move(*last_oper));
    }

    last_oper = &predicate_oper;
  }

  unique_ptr<LogicalOperator> group_by_oper;
  rc = create_group_by_plan(select_stmt, group_by_oper);
  if (OB_FAIL(rc)) {
    LOG_WARN("failed to create group by logical plan. rc=%s", strrc(rc));
    return rc;
  }

  if (group_by_oper) {
    if (*last_oper) {
      group_by_oper->add_child(std::move(*last_oper));
    }

    last_oper = &group_by_oper;
  }

  // 如果有group_by的话，应该再添加一层predicate算子
  unique_ptr<LogicalOperator> having_oper;
  if (group_by_oper != nullptr) {
    rc = create_plan(select_stmt->having(), having_oper);
    if (OB_FAIL(rc)) {
      LOG_WARN("failed to create having logical plan. rc=%s", strrc(rc));
      return rc;
    }
    if (having_oper) {
      if (*last_oper) {
        having_oper->add_child(std::move(*last_oper));
      }

      last_oper = &having_oper;
    }
  }

  auto project_oper = make_unique<ProjectLogicalOperator>(std::move(select_stmt->query_expressions()));
  if (*last_oper) {
    project_oper->add_child(std::move(*last_oper));
  }
  project_oper->set_order_by(std::move(select_stmt->order_by()));

  logical_operator = std::move(project_oper);
  return RC::SUCCESS;
}

RC LogicalPlanGenerator::create_plan(FilterStmt *filter_stmt, unique_ptr<LogicalOperator> &logical_operator)
{
  RC                                  rc = RC::SUCCESS;
  std::vector<unique_ptr<Expression>> cmp_exprs;
  const std::vector<FilterUnit *>    &filter_units = filter_stmt->filter_units();
  for (FilterUnit *filter_unit : filter_units) {
    FilterObj &&filter_obj_left = filter_unit->left();
    // 不使用const是因为可能右边有子查询，需要std::move
    FilterObj &&filter_obj_right = filter_unit->right();

    if (filter_obj_right.is_sub_query() &&
        (filter_unit->comp() == CompOp::EXISTS_OP || filter_unit->comp() == CompOp::NOT_EXISTS_OP)) {
      // 这里是exists的情况，没有左值，只有右边一个子查询
      unique_ptr<Expression> left;
      unique_ptr<Expression> right = std::move(filter_obj_right.simple_sub_query);

      // 下面把right中的stmt转变为逻辑节点
      SubQueryExpr               *sub_query_expr = static_cast<SubQueryExpr *>(right.get());
      Stmt                       *stmt           = sub_query_expr->get_select_stmt().get();
      unique_ptr<LogicalOperator> sub_query_logical_operator;
      LogicalPlanGenerator::create(stmt, sub_query_logical_operator);
      sub_query_expr->set_logical_operator(std::move(sub_query_logical_operator));

      ComparisonExpr *cmp_expr = new ComparisonExpr(filter_unit->comp(), std::move(left), std::move(right));
      cmp_expr->set_conjunction_type(filter_unit->conjunction_type());
      cmp_exprs.emplace_back(cmp_expr);
      continue;
    }

    if (filter_obj_right.is_tuple_values()) {
      // 左边是属性，右边是一个直接的元组
      unique_ptr<Expression> left  = std::move(filter_obj_left.obj_expr);
      unique_ptr<Expression> right = std::move(filter_obj_right.tuple_values);

      ComparisonExpr *cmp_expr = new ComparisonExpr(filter_unit->comp(), std::move(left), std::move(right));
      cmp_expr->set_conjunction_type(filter_unit->conjunction_type());
      cmp_exprs.emplace_back(cmp_expr);
      continue;
    }

    if (filter_obj_right.is_sub_query() && !filter_obj_left.is_sub_query()) {
      // 左边是属性，右边是子查询
      unique_ptr<Expression> left  = std::move(filter_obj_left.obj_expr);
      unique_ptr<Expression> right = std::move(filter_obj_right.simple_sub_query);

      // 下面把right中的stmt转变为逻辑节点
      SubQueryExpr               *sub_query_expr = static_cast<SubQueryExpr *>(right.get());
      Stmt                       *stmt           = sub_query_expr->get_select_stmt().get();
      unique_ptr<LogicalOperator> sub_query_logical_operator;
      LogicalPlanGenerator::create(stmt, sub_query_logical_operator);
      sub_query_expr->set_logical_operator(std::move(sub_query_logical_operator));

      ComparisonExpr *cmp_expr = new ComparisonExpr(filter_unit->comp(), std::move(left), std::move(right));
      cmp_expr->set_conjunction_type(filter_unit->conjunction_type());
      cmp_exprs.emplace_back(cmp_expr);
      continue;
    }

    if (filter_obj_left.is_sub_query() && !filter_obj_right.is_sub_query()) {
      // 右边是属性，左边是子查询
      unique_ptr<Expression> right = std::move(filter_obj_right.obj_expr);
      unique_ptr<Expression> left  = std::move(filter_obj_left.simple_sub_query);

      // 下面把left中的stmt转变为逻辑节点
      SubQueryExpr               *sub_query_expr = static_cast<SubQueryExpr *>(left.get());
      Stmt                       *stmt           = sub_query_expr->get_select_stmt().get();
      unique_ptr<LogicalOperator> sub_query_logical_operator;
      LogicalPlanGenerator::create(stmt, sub_query_logical_operator);
      sub_query_expr->set_logical_operator(std::move(sub_query_logical_operator));

      ComparisonExpr *cmp_expr = new ComparisonExpr(filter_unit->comp(), std::move(left), std::move(right));
      cmp_expr->set_conjunction_type(filter_unit->conjunction_type());
      cmp_exprs.emplace_back(cmp_expr);
      continue;
    }

    if (filter_obj_left.is_sub_query() && filter_obj_right.is_sub_query()) {
      // 左右两边都是子查询
      unique_ptr<Expression> left  = std::move(filter_obj_left.simple_sub_query);
      unique_ptr<Expression> right = std::move(filter_obj_right.simple_sub_query);

      // 下面把left中的stmt转变为逻辑节点
      SubQueryExpr               *left_sub_query_expr = static_cast<SubQueryExpr *>(left.get());
      Stmt                       *left_stmt           = left_sub_query_expr->get_select_stmt().get();
      unique_ptr<LogicalOperator> left_sub_query_logical_operator;
      LogicalPlanGenerator::create(left_stmt, left_sub_query_logical_operator);
      left_sub_query_expr->set_logical_operator(std::move(left_sub_query_logical_operator));

      // 下面把right中的stmt转变为逻辑节点
      SubQueryExpr               *right_sub_query_expr = static_cast<SubQueryExpr *>(right.get());
      Stmt                       *right_stmt           = right_sub_query_expr->get_select_stmt().get();
      unique_ptr<LogicalOperator> right_sub_query_logical_operator;
      LogicalPlanGenerator::create(right_stmt, right_sub_query_logical_operator);
      right_sub_query_expr->set_logical_operator(std::move(right_sub_query_logical_operator));

      ComparisonExpr *cmp_expr = new ComparisonExpr(filter_unit->comp(), std::move(left), std::move(right));
      cmp_expr->set_conjunction_type(filter_unit->conjunction_type());
      cmp_exprs.emplace_back(cmp_expr);
      continue;
    }

    // 现在不管是属性还是值，都变成Expression了

    unique_ptr<Expression> left  = std::move(filter_obj_left.obj_expr);
    unique_ptr<Expression> right = std::move(filter_obj_right.obj_expr);

    // 这里需要判断有没有NULL，如果有的话这里就不需要做类型转换了
    // 暂时不去考虑如果是一个表达式树的情况
    if (left->type() == ExprType::VALUE && left->type() != ExprType::ARITHMETIC) {
      Value left_val;
      if (OB_FAIL(rc = left->try_get_value(left_val))) {
        LOG_WARN("failed to get value from left child", strrc(rc));
        return rc;
      }
      if (left_val.is_null()) {
        goto lable;
      }
    }
    if (right->type() == ExprType::VALUE && right->type() != ExprType::ARITHMETIC) {
      Value right_val;
      if (OB_FAIL(rc = right->try_get_value(right_val))) {
        LOG_WARN("failed to get value from right child", strrc(rc));
        return rc;
      }
      if (right_val.is_null()) {
        goto lable;
      }
    }

    if (left->value_type() != right->value_type()) {
      auto left_to_right_cost = implicit_cast_cost(left->value_type(), right->value_type());
      auto right_to_left_cost = implicit_cast_cost(right->value_type(), left->value_type());
      if (left_to_right_cost <= right_to_left_cost && left_to_right_cost != INT32_MAX) {
        ExprType left_type = left->type();
        auto     cast_expr = make_unique<CastExpr>(std::move(left), right->value_type());
        if (left_type == ExprType::VALUE) {
          Value left_val;
          if (OB_FAIL(rc = cast_expr->try_get_value(left_val))) {
            LOG_WARN("failed to get value from left child", strrc(rc));
            return rc;
          }
          left = make_unique<ValueExpr>(left_val);
        } else {
          left = std::move(cast_expr);
        }
      } else if (right_to_left_cost < left_to_right_cost && right_to_left_cost != INT32_MAX) {
        ExprType right_type = right->type();
        auto     cast_expr  = make_unique<CastExpr>(std::move(right), left->value_type());
        if (right_type == ExprType::VALUE) {
          Value right_val;
          if (OB_FAIL(rc = cast_expr->try_get_value(right_val))) {
            LOG_WARN("failed to get value from right child", strrc(rc));
            return rc;
          }
          right = make_unique<ValueExpr>(right_val);
        } else {
          right = std::move(cast_expr);
        }

      } else {
        rc = RC::UNSUPPORTED;
        LOG_WARN("unsupported cast from %s to %s", attr_type_to_string(left->value_type()), attr_type_to_string(right->value_type()));
        return rc;
      }
    }

  lable:

    ComparisonExpr *cmp_expr = new ComparisonExpr(filter_unit->comp(), std::move(left), std::move(right));
    cmp_expr->set_conjunction_type(filter_unit->conjunction_type());
    cmp_exprs.emplace_back(cmp_expr);
  }

  unique_ptr<PredicateLogicalOperator> predicate_oper;
  // 这里最好是做一个嵌套，使用多层Conjuncion嵌套的形式
  // 只要后面是AND,那么统统放在一个Conjunction中，遇到OR,新建一个Conjuncion,然后嵌套进去
  // 最后再把OR的放在一个最外层的Conjunction中，这是因为AND的优先级高于OR,要先结合AND
  vector<unique_ptr<Expression>> conjunction_exprs;
  unsigned                       i = 0;
  while (i < cmp_exprs.size()) {

    vector<unique_ptr<Expression>> and_children;
    for (; i < cmp_exprs.size(); i++) {
      ComparisonExpr *cmp_expr_tmp = static_cast<ComparisonExpr *>(cmp_exprs[i].get());
      if (cmp_expr_tmp->conjunction_type() == 0) {
        and_children.push_back(std::move(cmp_exprs[i]));
        i++;
        break;
      } else {
        and_children.push_back(std::move(cmp_exprs[i]));
      }
    }
    if (!and_children.empty()) {
      unique_ptr<ConjunctionExpr> and_expr(new ConjunctionExpr(ConjunctionExpr::Type::AND, and_children));
      conjunction_exprs.push_back(std::move(and_expr));
    }
  }
  if (!conjunction_exprs.empty()) {
    unique_ptr<ConjunctionExpr> conjunction_expr(new ConjunctionExpr(ConjunctionExpr::Type::OR, conjunction_exprs));
    predicate_oper = unique_ptr<PredicateLogicalOperator>(new PredicateLogicalOperator(std::move(conjunction_expr)));
  }

  logical_operator = std::move(predicate_oper);
  return rc;
}

int LogicalPlanGenerator::implicit_cast_cost(AttrType from, AttrType to)
{
  if (from == to) {
    return 0;
  }
  return DataType::type_instance(from)->cast_cost(to);
}

RC LogicalPlanGenerator::create_plan(InsertStmt *insert_stmt, unique_ptr<LogicalOperator> &logical_operator)
{
  Table        *table = insert_stmt->table();
  vector<Value> values(insert_stmt->values(), insert_stmt->values() + insert_stmt->value_amount());

  InsertLogicalOperator *insert_operator = new InsertLogicalOperator(table, values);
  logical_operator.reset(insert_operator);
  return RC::SUCCESS;
}

RC LogicalPlanGenerator::create_plan(UpdateStmt *update_stmt, unique_ptr<LogicalOperator> &logical_operator)
{
  Table *table = update_stmt->table();
  // vector<Value *> value = update_stmt->value();
  vector<FieldMeta *>         field_meta  = update_stmt->field_meta();
  FilterStmt                 *filter_stmt = update_stmt->filter_stmt();
  unique_ptr<LogicalOperator> table_get_oper(new TableGetLogicalOperator(table, ReadWriteMode::READ_WRITE));

  unique_ptr<LogicalOperator> predicate_oper;

  RC rc = create_plan(filter_stmt, predicate_oper);
  if (rc != RC::SUCCESS) {
    return rc;
  }

  // TODO: 这里应该对每一个表达式处理一下，对Value和子查询不同处理
  for (auto &expr : update_stmt->value()) {
    if (expr->type() == ExprType::VALUE) {
      continue;
    } else if (expr->type() == ExprType::SUB_QUERY) {
      SubQueryExpr                    *sub_query_expr = static_cast<SubQueryExpr *>(expr.get());
      std::unique_ptr<LogicalOperator> sub_query_logical_operator;
      RC rc2 = create_plan(sub_query_expr->get_select_stmt().get(), sub_query_logical_operator);
      if (rc2 != RC::SUCCESS) {
        LOG_WARN("failed to create logical plan for sub query");
        return RC::INTERNAL;
      }
      // 将stmt转换为logical_operator
      sub_query_expr->set_logical_operator(std::move(sub_query_logical_operator));
    } else {
      LOG_WARN("INVALID_ARGUMENT!");
      return RC::INVALID_ARGUMENT;
    }
  }
  unique_ptr<LogicalOperator> update_oper(
      new UpdateLogicalOperator(table, std::move(update_stmt->value()), field_meta));
  if (predicate_oper) {
    predicate_oper->add_child(std::move(table_get_oper));
    update_oper->add_child(std::move(predicate_oper));
  } else {
    update_oper->add_child(std::move(table_get_oper));
  }
  logical_operator = std::move(update_oper);
  return RC::SUCCESS;
}

RC LogicalPlanGenerator::create_plan(DeleteStmt *delete_stmt, unique_ptr<LogicalOperator> &logical_operator)
{
  Table                      *table       = delete_stmt->table();
  FilterStmt                 *filter_stmt = delete_stmt->filter_stmt();
  unique_ptr<LogicalOperator> table_get_oper(new TableGetLogicalOperator(table, ReadWriteMode::READ_WRITE));

  unique_ptr<LogicalOperator> predicate_oper;

  RC rc = create_plan(filter_stmt, predicate_oper);
  if (rc != RC::SUCCESS) {
    return rc;
  }

  unique_ptr<LogicalOperator> delete_oper(new DeleteLogicalOperator(table));

  if (predicate_oper) {
    predicate_oper->add_child(std::move(table_get_oper));
    delete_oper->add_child(std::move(predicate_oper));
  } else {
    delete_oper->add_child(std::move(table_get_oper));
  }

  logical_operator = std::move(delete_oper);
  return rc;
}

RC LogicalPlanGenerator::create_plan(ExplainStmt *explain_stmt, unique_ptr<LogicalOperator> &logical_operator)
{
  unique_ptr<LogicalOperator> child_oper;

  Stmt *child_stmt = explain_stmt->child();

  RC rc = create(child_stmt, child_oper);
  if (rc != RC::SUCCESS) {
    LOG_WARN("failed to create explain's child operator. rc=%s", strrc(rc));
    return rc;
  }

  logical_operator = unique_ptr<LogicalOperator>(new ExplainLogicalOperator);
  logical_operator->add_child(std::move(child_oper));
  return rc;
}

RC LogicalPlanGenerator::create_group_by_plan(SelectStmt *select_stmt, unique_ptr<LogicalOperator> &logical_operator)
{
  vector<unique_ptr<Expression>>             &group_by_expressions = select_stmt->group_by();
  vector<Expression *>                        aggregate_expressions;
  vector<unique_ptr<Expression>>             &query_expressions = select_stmt->query_expressions();
  function<RC(std::unique_ptr<Expression> &)> collector         = [&](unique_ptr<Expression> &expr) -> RC {
    RC rc = RC::SUCCESS;
    if (expr->type() == ExprType::AGGREGATION) {
      expr->set_pos(aggregate_expressions.size() + group_by_expressions.size());
      aggregate_expressions.push_back(expr.get());
    }
    rc = ExpressionIterator::iterate_child_expr(*expr, collector);
    return rc;
  };

  function<RC(std::unique_ptr<Expression> &)> bind_group_by_expr = [&](unique_ptr<Expression> &expr) -> RC {
    RC rc = RC::SUCCESS;
    for (size_t i = 0; i < group_by_expressions.size(); i++) {
      auto &group_by = group_by_expressions[i];
      if (expr->type() == ExprType::AGGREGATION) {
        break;
      } else if (expr->equal(*group_by)) {
        expr->set_pos(i);
        continue;
      } else {
        rc = ExpressionIterator::iterate_child_expr(*expr, bind_group_by_expr);
      }
    }
    return rc;
  };

  bool                                        found_unbound_column = false;
  function<RC(std::unique_ptr<Expression> &)> find_unbound_column  = [&](unique_ptr<Expression> &expr) -> RC {
    RC rc = RC::SUCCESS;
    if (expr->type() == ExprType::AGGREGATION) {
      // do nothing
    } else if (expr->pos() != -1) {
      // do nothing
    } else if (expr->type() == ExprType::FIELD) {
      found_unbound_column = true;
    } else {
      rc = ExpressionIterator::iterate_child_expr(*expr, find_unbound_column);
    }
    return rc;
  };

  for (unique_ptr<Expression> &expression : query_expressions) {
    bind_group_by_expr(expression);
  }

  for (unique_ptr<Expression> &expression : query_expressions) {
    find_unbound_column(expression);
  }

  // collect all aggregate expressions
  for (unique_ptr<Expression> &expression : query_expressions) {
    collector(expression);
  }

  if (group_by_expressions.empty() && aggregate_expressions.empty()) {
    // 既没有group by也没有聚合函数，不需要group by
    return RC::SUCCESS;
  }

  if (found_unbound_column) {
    LOG_WARN("column must appear in the GROUP BY clause or must be part of an aggregate function");
    return RC::INVALID_ARGUMENT;
  }

  // 如果只需要聚合，但是没有group by 语句，需要生成一个空的group by 语句

  auto group_by_oper =
      make_unique<GroupByLogicalOperator>(std::move(group_by_expressions), std::move(aggregate_expressions));
  logical_operator = std::move(group_by_oper);
  return RC::SUCCESS;
}
